#!/bin/bash

# 网络名称,同目录名称,需要模型审视修改
Network="BEVFusion"
export PERFORMANCE_MODE=1 # Performance-Testing mode

batch_size=4                    # 单卡batch_size
num_npu=8                       # 每节点NPU卡数
nnodes=2                        # 节点总数
node_rank=0                     # 当前节点编号（0 ~ nnodes-1），默认为主节点
port=29500                      # 通信端口号
master_addr=""                  # 主节点IP地址
world_size=$((nnodes * num_npu))

# cd到与test文件夹同层级目录下执行脚本，提高兼容性；test_path_dir为包含test文件夹的路径
cur_path=$(pwd)
cur_path_last_dirname=${cur_path##*/}
if [ x"${cur_path_last_dirname}" == x"test" ]; then
  test_path_dir=${cur_path}
  cd ..
  cur_path=$(pwd)
else
  test_path_dir=${cur_path}/test
fi

source ${test_path_dir}/env_npu.sh
# 解析参数
source ${test_path_dir}/parse_args.sh
declare_required_params batch_size num_npu nnodes node_rank port master_addr # 接收参数顺序
parse_common_args "$@"

base_batch_size=$(($batch_size * $num_npu))

#创建DeviceID输出目录，不需要修改
output_path=${cur_path}/test/output/

mkdir -p ${output_path}

# 性能测试1个epoch
cd mmdetection3d
sed -i "s|max_epochs=6|max_epochs=1|g" projects/BEVFusion/configs/bevfusion_lidar-cam_voxel0075_second_secfpn_8xb4-cyclic-20e_nus-3d.py

#训练开始时间，不需要修改
start_time=$(date +%s)
bash tools/nnodes_dist_train.sh \
    projects/BEVFusion/configs/bevfusion_lidar-cam_voxel0075_second_secfpn_8xb4-cyclic-20e_nus-3d.py ${num_npu} ${nnodes} ${node_rank} ${port} ${master_addr}\
    --cfg-options train_dataloader.batch_size=${batch_size} auto_scale_lr.base_batch_size=${base_batch_size} \
    load_from=pretrained/bevfusion_lidar_voxel0075_second_secfpn_8xb4-cyclic-20e_nus-3d-2628f933.pth model.img_backbone.init_cfg.checkpoint=pretrained/swint-nuimages-pretrained.pth \
    > ${test_path_dir}/output/train_performance_${num_npu}p_base_fp32.log 2>&1 &

wait
#训练结束时间，不需要修改
end_time=$(date +%s)
e2e_time=$(($end_time - $start_time))

# 回退epoch数目
sed -i "s|max_epochs=1|max_epochs=6|g" projects/BEVFusion/configs/bevfusion_lidar-cam_voxel0075_second_secfpn_8xb4-cyclic-20e_nus-3d.py
cd ..

# 主节点打印性能数据
if [ "$node_rank" -eq 0 ]; then
  #结果打印，不需要修改
  echo "------------------ Final result ------------------"

  #获取性能数据，不需要修改
  #单迭代训练时长，不需要修改
  TrainingTime=$(grep -v val ${test_path_dir}/output/train_performance_${num_npu}p_base_fp32.log | grep -o " time: [0-9.]*"  | tail -n +200 | grep -o "[0-9.]*" | awk '{sum += $1} END {print sum/NR}')

  #吞吐量
  ActualFPS=$(awk BEGIN'{print ('$batch_size' * '$world_size') / '$TrainingTime'}')

  #打印，不需要修改
  echo "Final Performance images/sec : $ActualFPS"

  #loss值，不需要修改
  ActualLoss=$(grep -o "loss: [0-9.]*" ${test_path_dir}/output/train_performance_${num_npu}p_base_fp32.log | awk 'END {print $NF}')

  #打印，不需要修改
  echo "Final Train Loss : ${ActualLoss}"
  echo "E2E Training Duration sec : $e2e_time"

  #性能看护结果汇总
  #训练用例信息，不需要修改
  BatchSize=${batch_size}
  WORLD_SIZE=${world_size}
  DeviceType=$(uname -m)
  CaseName=${Network}_bs${BatchSize}_${WORLD_SIZE}'p'_'performance'

  #关键信息打印到${CaseName}.log中，不需要修改
  echo "Network = ${Network}" >${test_path_dir}/output/${CaseName}.log
  echo "RankSize = ${WORLD_SIZE}" >>${test_path_dir}/output/${CaseName}.log
  echo "BatchSize = ${BatchSize}" >>${test_path_dir}/output/${CaseName}.log
  echo "DeviceType = ${DeviceType}" >>${test_path_dir}/output/${CaseName}.log
  echo "CaseName = ${CaseName}" >>${test_path_dir}/output/${CaseName}.log
  echo "ActualFPS = ${ActualFPS}" >>${test_path_dir}/output/${CaseName}.log
  echo "TrainingTime = ${TrainingTime}" >>${test_path_dir}/output/${CaseName}.log
  echo "ActualLoss = ${ActualLoss}" >>${test_path_dir}/output/${CaseName}.log
  echo "E2ETrainingTime = ${e2e_time}" >>${test_path_dir}/output/${CaseName}.log
fi